#ifndef AST_H
#define AST_H

#include "token.h"
#include <vector>
#include <memory>
#include <string>

// Forward declarations
class ASTNode;
class Expression;
class Statement;
class Declaration;
class FunctionDeclaration;
class Program;

// Base class for all AST nodes
class ASTNode {
public:
    virtual ~ASTNode() = default;
    virtual std::string toString() const = 0;
};

// Expression types
class Expression : public ASTNode {
public:
    TokenType type;
    std::string value;
    
    Expression(TokenType type, const std::string& value) 
        : type(type), value(value) {}
    
    std::string toString() const override;
};

class BinaryExpression : public Expression {
public:
    std::unique_ptr<Expression> left;
    std::unique_ptr<Expression> right;
    TokenType operatorType;
    
    BinaryExpression(std::unique_ptr<Expression> left, TokenType op, std::unique_ptr<Expression> right)
        : Expression(TokenType::UNKNOWN, ""), left(std::move(left)), right(std::move(right)), operatorType(op) {}
    
    std::string toString() const override;
};

class UnaryExpression : public Expression {
public:
    std::unique_ptr<Expression> operand;
    TokenType operatorType;
    
    UnaryExpression(TokenType op, std::unique_ptr<Expression> operand)
        : Expression(TokenType::UNKNOWN, ""), operand(std::move(operand)), operatorType(op) {}
    
    std::string toString() const override;
};

class FunctionCall : public Expression {
public:
    std::string functionName;
    std::vector<std::unique_ptr<Expression>> arguments;
    
    FunctionCall(const std::string& name, std::vector<std::unique_ptr<Expression>> args)
        : Expression(TokenType::UNKNOWN, ""), functionName(name), arguments(std::move(args)) {}
    
    std::string toString() const override;
};

// Statement types
class Statement : public ASTNode {
public:
    virtual ~Statement() = default;
};

class ExpressionStatement : public Statement {
public:
    std::unique_ptr<Expression> expression;
    
    ExpressionStatement(std::unique_ptr<Expression> expr) 
        : expression(std::move(expr)) {}
    
    std::string toString() const override;
};

class ReturnStatement : public Statement {
public:
    std::unique_ptr<Expression> expression;
    
    ReturnStatement(std::unique_ptr<Expression> expr = nullptr) 
        : expression(std::move(expr)) {}
    
    std::string toString() const override;
};

class IfStatement : public Statement {
public:
    std::unique_ptr<Expression> condition;
    std::unique_ptr<Statement> thenBranch;
    std::unique_ptr<Statement> elseBranch;
    
    IfStatement(std::unique_ptr<Expression> cond, std::unique_ptr<Statement> thenStmt, 
                std::unique_ptr<Statement> elseStmt = nullptr)
        : condition(std::move(cond)), thenBranch(std::move(thenStmt)), elseBranch(std::move(elseStmt)) {}
    
    std::string toString() const override;
};

class WhileStatement : public Statement {
public:
    std::unique_ptr<Expression> condition;
    std::unique_ptr<Statement> body;
    
    WhileStatement(std::unique_ptr<Expression> cond, std::unique_ptr<Statement> body)
        : condition(std::move(cond)), body(std::move(body)) {}
    
    std::string toString() const override;
};

class BlockStatement : public Statement {
public:
    std::vector<std::unique_ptr<Statement>> statements;
    
    BlockStatement(std::vector<std::unique_ptr<Statement>> stmts)
        : statements(std::move(stmts)) {}
    
    std::string toString() const override;
};

// Declaration types
class Declaration : public ASTNode {
public:
    virtual ~Declaration() = default;
};

class VariableDeclaration : public Declaration {
public:
    TokenType type;
    std::string name;
    std::unique_ptr<Expression> initializer;
    
    VariableDeclaration(TokenType type, const std::string& name, std::unique_ptr<Expression> init = nullptr)
        : type(type), name(name), initializer(std::move(init)) {}
    
    std::string toString() const override;
};

class FunctionDeclaration : public Declaration {
public:
    TokenType returnType;
    std::string name;
    std::vector<std::pair<TokenType, std::string>> parameters;
    std::unique_ptr<BlockStatement> body;
    
    FunctionDeclaration(TokenType returnType, const std::string& name, 
                       std::vector<std::pair<TokenType, std::string>> params,
                       std::unique_ptr<BlockStatement> body)
        : returnType(returnType), name(name), parameters(std::move(params)), body(std::move(body)) {}
    
    std::string toString() const override;
};

// Program root
class Program : public ASTNode {
public:
    std::vector<std::unique_ptr<Declaration>> declarations;
    
    Program(std::vector<std::unique_ptr<Declaration>> decls)
        : declarations(std::move(decls)) {}
    
    std::string toString() const override;
};

#endif
